{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9HcmLuWtE213"
      },
      "source": [
        "# Finetuning Magenta RT\n",
        "\n",
        "<a href=\"https://colab.research.google.com/github/magenta/magenta-realtime/blob/main/notebooks/Magenta_RT_Finetune.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
        "\n",
        "This notebook shows how [Magenta RealTime](https://g.co/magenta/rt) can be customized through finetuning. You can either (1) follow the steps to train on your own audio dataset or (2) explore an existing finetuned model.\n",
        "\n",
        "### 1. Finetuning Magenta RT on your own data\n",
        "For finetuning, follow the notebook step by step. We recommend at least ~30 minutes of audio of consistent style, but you're free to experiment with something different. Finetuning allows you to build a custom model that sounds uniquely yours, so curating your own training data, and experimenting with it, is an important part of the process.\n",
        "\n",
        "### 2. Exploring finetuned models\n",
        "#### ⭐️ Holly Herndon x Magenta RT ⭐️\n",
        "We're thrilled to have worked with [Holly Herndon](https://herndondryhurst.studio/) to finetune a model on her vocal style to share with all of you. To play with this model, run [Step 1](#scrollTo=RP-rwG3Uzz_1) and [Step 4](#scrollTo=94S5nEXipBAY) of the notebook, selecting `holly_finetune` as the checkpoint to load.\n",
        "\n",
        "Compared to the [base Magenta RT demo](https://colab.research.google.com/github/magenta/magenta-realtime/blob/main/notebooks/Magenta_RT_Demo.ipynb), this demo includes additional prompt controls that allow you to steer the model towards styles that are more in-distribution with respect to the finetuning dataset.\n",
        "\n",
        "We particularly recommend experimenting with the in-distribution controls with this model and noticing how mixing in the mean and centroids affects the generation. You might hear vocalizations that contour many the generated sounds, or, in some cases, \"color\" the instruments with a timbre that resembles Holly's voice (e.g. trumpet/violin oscillate between what you'd expect and her voice).\n",
        "\n",
        "#### Loading your own finetuned model\n",
        "If you have already finetuned your own model in a previous session and would like to load it in the demo, run [Step 1](#scrollTo=RP-rwG3Uzz_1) and follow the instructions in [Step 4](#scrollTo=94S5nEXipBAY)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RP-rwG3Uzz_1"
      },
      "source": [
        "# Step 1: 😴 One-time setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qJ4kA_ZEtdsA",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title **Run this cell** to install dependencies (~5 minutes)\n",
        "# @markdown Make sure you are running on **`v5e-1 TPU` runtime** via `Runtime > Change Runtime Type`\n",
        "\n",
        "# @markdown Colab may prompt you to restart session. **Wait until the cell finishes running to restart**!\n",
        "\n",
        "# Clone library\n",
        "!git clone https://github.com/magenta/magenta-realtime.git\n",
        "!git clone https://github.com/google-research/t5x.git\n",
        "\n",
        "# Install library and dependencies\n",
        "# If running on TPU (recommended, runs on free tier Colab TPUs):\n",
        "!pip install -e t5x/[tpu] && pip install -e magenta-realtime/[tpu] && pip install tf2jax==0.3.8\n",
        "\n",
        "# Uncomment if running on GPU (requires A100 via Colab Pro to be fast enough):\n",
        "# !patch t5x/setup.py < magenta-realtime/patch/t5x_setup.py.patch\n",
        "# !patch t5x/t5x/partitioning.py < magenta-realtime/patch/t5x_partitioning.py.patch\n",
        "# !pip install -e t5x/[gpu] && pip install -e magenta-realtime/[gpu] && pip install tf2jax==0.3.8\n",
        "\n",
        "!sed -i '/import tensorflow_text as tf_text/d' /usr/local/lib/python3.12/dist-packages/seqio/vocabularies.py\n",
        "!sed -i \"s|device_kind == 'TPU v4 lite'|device_kind == 'TPU v4 lite' or device_kind == 'TPU v5 lite'|g\" /usr/local/lib/python3.12/dist-packages/t5x/partitioning.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "15YXTfWz2Q35"
      },
      "source": [
        "# Step 2: 📁 Prepare the training data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "763S1sRn2Rfv",
        "language": "python"
      },
      "outputs": [],
      "source": [
        "# @title Run this cell to load and process the training data\n",
        "\n",
        "# @markdown Upload your training data to a Google Drive folder or directly to Colab using the file browser on the left, and specify the name of the directory below in `AUDIO_FOLDER_NAME`.\n",
        "# @markdown For example, if you put your audio in a folder called \"Guitar\" in the root directory on Google Drive, select `AUDIO_SOURCE: \"drive\"` and set `AUDIO_FOLDER_NAME: \"Guitar\"`. You can also reference subdirectories on Google Drive, e.g., `AUDIO_FOLDER_NAME: \"MyAudio/Guitar\"`.\n",
        "\n",
        "AUDIO_SOURCE = \"drive\" # @param [\"colab\",\"drive\"]\n",
        "AUDIO_FOLDER_NAME = \"\"  #@param {type:\"string\", \"placeholder\": \"Name of the top-level folder containing your audio data\"}\n",
        "AUDIO_EXTENSIONS = \"wav,mp3,flac,ogg\"  #@param {type:\"string\"}\n",
        "FILTER_QUIET_AUDIO = False # @param {type:\"boolean\"}\n",
        "\n",
        "import os\n",
        "import seqio\n",
        "import pathlib\n",
        "import numpy as np\n",
        "import t5x\n",
        "import clu.data\n",
        "import tensorflow as tf\n",
        "import tensorflow.data as tf_data\n",
        "import tensorflow.io as tf_io\n",
        "from matplotlib import pyplot as plt\n",
        "from google.colab import drive\n",
        "from tqdm.notebook import tqdm\n",
        "from sklearn.cluster import KMeans\n",
        "from IPython.display import display, Audio\n",
        "\n",
        "from magenta_rt.finetune import data\n",
        "from magenta_rt.finetune import tasks\n",
        "from magenta_rt import audio as audio_lib\n",
        "\n",
        "\n",
        "if AUDIO_SOURCE == \"drive\":\n",
        "  drive.mount('/content/drive', force_remount=True)\n",
        "  audio_dir = f'/content/drive/MyDrive/{AUDIO_FOLDER_NAME}'\n",
        "else:\n",
        "  audio_dir = f'/content/{AUDIO_FOLDER_NAME}'\n",
        "\n",
        "AUDIO_DIR = pathlib.Path(audio_dir)\n",
        "if not AUDIO_DIR.is_dir():\n",
        "  raise FileNotFoundError(f\"Audio directory {audio_dir} does not exist\")\n",
        "\n",
        "# Find audio paths\n",
        "print(f\"Searching {AUDIO_DIR} for audio files...\")\n",
        "audio_extensions = [e.strip() for e in AUDIO_EXTENSIONS.split(',') if len(e.strip()) > 0]\n",
        "if len(audio_extensions) == 0:\n",
        "  raise ValueError(\"No audio extensions specified\")\n",
        "AUDIO_PATHS = []\n",
        "for e in audio_extensions:\n",
        "  AUDIO_PATHS.extend(list(AUDIO_DIR.glob(f'**/*.{e}')))\n",
        "AUDIO_PATHS = sorted(AUDIO_PATHS)\n",
        "\n",
        "# Help user\n",
        "if len(AUDIO_PATHS) == 0:\n",
        "  raise FileNotFoundError(f\"No audio files found in {audio_dir} with extensions {audio_extensions}\")\n",
        "else:\n",
        "  print(f\"Found {len(AUDIO_PATHS)} audio files in {audio_dir}. A few examples:\")\n",
        "  for p in AUDIO_PATHS[:5]:\n",
        "    print('-' * 80)\n",
        "    print(p)\n",
        "    display(Audio(p))\n",
        "  print('...')\n",
        "\n",
        "\n",
        "TASK_NAME = str(AUDIO_DIR.stem)\n",
        "OUTPUT_DIR = str(pathlib.Path(pathlib.Path.cwd() / 'mrt_finetune'))\n",
        "if not os.path.exists(OUTPUT_DIR):\n",
        "  os.makedirs(OUTPUT_DIR)\n",
        "OUTPUT_PATTERN = f'{OUTPUT_DIR}/{TASK_NAME}_examples.recordio'\n",
        "\n",
        "\n",
        "print(\"Tokenizing the training data...\")\n",
        "featurizer = data.Featurizer(\n",
        "    filter_quiet=FILTER_QUIET_AUDIO,\n",
        "    min_clip_seconds=2,\n",
        "    include_style_embeddings=True,\n",
        ")\n",
        "\n",
        "records_count = 0\n",
        "with tf_io.TFRecordWriter(OUTPUT_PATTERN) as file_writer:\n",
        "  for audio_path in tqdm(AUDIO_PATHS):\n",
        "    audio_input = audio_lib.Waveform.from_file(audio_path)\n",
        "    tokenized_iter = featurizer.process(audio_input)\n",
        "    for tokenized_example in tokenized_iter:\n",
        "      records_count += 1\n",
        "      file_writer.write(tokenized_example.SerializeToString())\n",
        "\n",
        "print(f'{records_count} records written')\n",
        "feaurized_audio_length = (records_count) * 30\n",
        "print(f'Total duration of featurized audio: {feaurized_audio_length:.0f} seconds ({(feaurized_audio_length/60):.1f} minutes)')\n",
        "\n",
        "\n",
        "print(f\"Registering new Seqio task...\")\n",
        "if TASK_NAME in seqio.TaskRegistry.names():\n",
        "  seqio.TaskRegistry.remove(TASK_NAME)\n",
        "  seqio.TaskRegistry.remove(TASK_NAME+\"_eval\")\n",
        "\n",
        "tasks.register_task(\n",
        "    name=TASK_NAME,\n",
        "    split_to_filepattern={\n",
        "        'train': OUTPUT_PATTERN,\n",
        "        'validation': OUTPUT_PATTERN,\n",
        "    },\n",
        "    reader_cls=tf_data.TFRecordDataset,\n",
        "    acoustic_key='acoustic_tokens',\n",
        "    style_key='style_tokens',\n",
        "    encoder_codec_rvq_depth=4,\n",
        "    decoder_codec_rvq_depth=16,\n",
        "    max_prompt_secs=10,\n",
        ")\n",
        "\n",
        "print(f\"{TASK_NAME} SeqIO task registered\")\n",
        "\n",
        "\n",
        "print(\"Computing statistics on the finetuning style embeddings...\")\n",
        "def decode_fn(record_bytes):\n",
        "  return tf.io.parse_single_example(\n",
        "      record_bytes,\n",
        "      {\"style_embeddings\": tf.io.FixedLenFeature([], dtype=tf.string)}\n",
        "  )\n",
        "\n",
        "audio_style_embeddings = []\n",
        "for batch in tf.data.TFRecordDataset([OUTPUT_PATTERN]).map(decode_fn):\n",
        "  style_embeds = tf.io.parse_tensor(batch['style_embeddings'], out_type=tf.float32).numpy()\n",
        "  audio_style_embeddings.append(np.mean(style_embeds, axis=0))\n",
        "audio_style_embeddings = np.array(audio_style_embeddings)\n",
        "np.save(f'{OUTPUT_DIR}/{TASK_NAME}_style_embeddings.npy', audio_style_embeddings)\n",
        "mean_style_embed = np.mean(audio_style_embeddings, axis=0)\n",
        "kmeans = KMeans(n_clusters=5, random_state=0, n_init=10)\n",
        "kmeans.fit(audio_style_embeddings)\n",
        "cluster_centroids = kmeans.cluster_centers_"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M0vfNy47ov2c"
      },
      "source": [
        "# Step 3: 🔧 Finetune"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "TsR531sXowou",
        "language": "python"
      },
      "outputs": [],
      "source": [
        "# @title Run this cell to start finetuning\n",
        "\n",
        "# @markdown If you'd like to save the experiment outputs (checkpoints,\n",
        "# @markdown preprocessed data, style embeddings for in-distribution steering) to Google Drive so you can load the\n",
        "# @markdown finetuned model in a future session tick the\n",
        "# @markdown `COPY_EXPERIMENT_TO_DRIVE` box before running this cell.\n",
        "\n",
        "from datetime import datetime\n",
        "\n",
        "from magenta_rt.finetune import finetuner\n",
        "\n",
        "EXPERIMENT_NAME = \"\" # @param {type: \"string\", \"placeholder\": \"Name of your finetuning experiment (optional)\"}\n",
        "MODEL_SIZE = 'large' # @param ['base', 'large']\n",
        "FINETUNING_STEPS = 6000 # @param\n",
        "SAVE_CKP_PERIOD = 1000 # @param\n",
        "COPY_EXPERIMENT_TO_DRIVE = False # @param {type:\"boolean\"}\n",
        "\n",
        "if not EXPERIMENT_NAME:\n",
        "  EXPERIMENT_NAME = datetime.now().strftime(\"%Y%m%d_%H%M\")\n",
        "MODEL_OUTPUT_DIR = f\"{OUTPUT_DIR}/{EXPERIMENT_NAME}\"\n",
        "\n",
        "def get_ds_iterator(\n",
        "    mixture_or_task_name: str,\n",
        "    batch_size: int = 8,\n",
        "    use_cached_tasks: bool = False,\n",
        "    shuffle=True,\n",
        "    ):\n",
        "  train_dataset_cfg = t5x.utils.DatasetConfig(\n",
        "      mixture_or_task_name = mixture_or_task_name,\n",
        "      task_feature_lengths = {'inputs': 1006, 'targets': 800},\n",
        "      split = 'train',\n",
        "      batch_size = batch_size,\n",
        "      shuffle = shuffle,\n",
        "      use_cached = use_cached_tasks,\n",
        "      pack = True,\n",
        "      module = None,\n",
        "      seed = 42,\n",
        "  )\n",
        "\n",
        "  train_ds = t5x.utils.get_dataset(\n",
        "      cfg=train_dataset_cfg,\n",
        "      shard_id=0,\n",
        "      num_shards=1,\n",
        "      feature_converter_cls=seqio.EncDecFeatureConverter,\n",
        "  )\n",
        "  train_iter = clu.data.dataset_iterator.TfDatasetIterator(train_ds, checkpoint=False)\n",
        "  return train_iter\n",
        "\n",
        "def plot_training_curves(training_summary, save=True):\n",
        "  num_plots = len(training_summary.keys())\n",
        "  fig, axs = plt.subplots(1, num_plots, figsize=(5*num_plots, 5))\n",
        "  fig.suptitle(f\"Training curves for {EXPERIMENT_NAME}\")\n",
        "  for i, (k, v) in enumerate(training_summary.items()):\n",
        "    axs[i].plot([i.value for i in v])\n",
        "    axs[i].set_xlabel('Step')\n",
        "    axs[i].set_ylabel(k)\n",
        "  if save:\n",
        "    plt.savefig(f\"{MODEL_OUTPUT_DIR}/training_curves.png\")\n",
        "\n",
        "print(\"Setting up the finetuner...\")\n",
        "MRTFinetuner = finetuner.MagentaRTFinetuner(\n",
        "    tag=MODEL_SIZE,\n",
        "    output_dir=MODEL_OUTPUT_DIR,\n",
        ")\n",
        "\n",
        "print(\"Training...\")\n",
        "train_iter = get_ds_iterator(TASK_NAME)\n",
        "\n",
        "MRTFinetuner.train(\n",
        "    train_iter=train_iter,\n",
        "    num_steps=FINETUNING_STEPS,\n",
        "    save_ckpt_period=SAVE_CKP_PERIOD,\n",
        "  )\n",
        "\n",
        "MRTFinetuner.train_summary\n",
        "\n",
        "plot_training_curves(\n",
        "    {\n",
        "        'Loss': MRTFinetuner.loss,\n",
        "        'Accuracy': MRTFinetuner.accuracy,\n",
        "    }\n",
        ")\n",
        "\n",
        "if COPY_EXPERIMENT_TO_DRIVE:\n",
        "  print(\"Copying experiment to Google Drive...\")\n",
        "  drive.mount('/content/drive', force_remount=True)\n",
        "  np.save(f'{MODEL_OUTPUT_DIR}/mean_style_embed.npy', mean_style_embed)\n",
        "  np.save(f'{MODEL_OUTPUT_DIR}/cluster_centroids.npy', cluster_centroids)\n",
        "  if not os.path.exists(f'/content/drive/MyDrive/{MODEL_OUTPUT_DIR}'):\n",
        "    ! cp -r {MODEL_OUTPUT_DIR} /content/drive/MyDrive/\n",
        "  else:\n",
        "    print(\"Experiment already exists in Google Drive\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "94S5nEXipBAY"
      },
      "source": [
        "# Step 4: 🎚 Play with the finetuned model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zYiU3yvJQZdR"
      },
      "outputs": [],
      "source": [
        "# @markdown If you'd like to load a model from a previous experiment,\n",
        "# @markdown specify path and model size below and run this cell.\n",
        "# @markdown You can load an experiment from Google Drive\n",
        "# @markdown (`EXPERIMENT_SOURCE = \"drive\"`) or directly from Colab\n",
        "# @markdown (`EXPERIMENT_SOURCE = \"colab\"`). If you choose colab,\n",
        "# @markdown you'll need to upload the experiment folder through the file browser\n",
        "# @markdown on the left.\n",
        "\n",
        "EXPERIMENT_SOURCE = \"colab\" # @param [\"colab\",\"drive\"]\n",
        "PREVIOUS_EXPERIMENT_FOLDER = \"\" # @param {type: \"string\", \"placeholder\": \"Type the path to the experiment folder here, e.g. mrt_finetune/my_experiment_name\"}\n",
        "FINETUNED_MODEL_SIZE = 'large' # @param ['base', 'large']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "w6l-_YV0pAHT"
      },
      "outputs": [],
      "source": [
        "# @title Select the checkpoint\n",
        "\n",
        "import os\n",
        "from google.colab import drive\n",
        "from ipywidgets import widgets\n",
        "from IPython.display import display\n",
        "\n",
        "try:\n",
        "  PREVIOUS_EXPERIMENT_FOLDER\n",
        "except:\n",
        "  PREVIOUS_EXPERIMENT_FOLDER = \"\"\n",
        "\n",
        "try:\n",
        "  MODEL_OUTPUT_DIR\n",
        "except:\n",
        "  MODEL_OUTPUT_DIR = \"\"\n",
        "\n",
        "\n",
        "if PREVIOUS_EXPERIMENT_FOLDER:\n",
        "  if EXPERIMENT_SOURCE == \"drive\":\n",
        "    drive.mount('/content/drive', force_remount=True)\n",
        "    MODEL_OUTPUT_DIR = f'/content/drive/MyDrive/{PREVIOUS_EXPERIMENT_FOLDER}'\n",
        "  else:\n",
        "    MODEL_OUTPUT_DIR = f'/content/{PREVIOUS_EXPERIMENT_FOLDER}'\n",
        "  MODEL_SIZE = FINETUNED_MODEL_SIZE\n",
        "  mean_style_embed = np.load(f\"{MODEL_OUTPUT_DIR}/mean_style_embed.npy\")\n",
        "  cluster_centroids = np.load(f\"{MODEL_OUTPUT_DIR}/cluster_centroids.npy\")\n",
        "\n",
        "if MODEL_OUTPUT_DIR:\n",
        "  checkpoints = [int(dir.split(\"_\")[1]) for dir in os.listdir(MODEL_OUTPUT_DIR) if dir.startswith(\"checkpoint_\")]\n",
        "  checkpoints.sort()\n",
        "else:\n",
        "  checkpoints = []\n",
        "checkpoints = [\"holly_finetune\"] + checkpoints\n",
        "\n",
        "checkpoint_to_load = widgets.Dropdown(\n",
        "    options=checkpoints,\n",
        "    value=checkpoints[0],\n",
        "    description='Checkpoint:',\n",
        "    disabled=False,\n",
        ")\n",
        "\n",
        "display(checkpoint_to_load)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LvXU2Il4WkBb"
      },
      "source": [
        "**Run the cell below and click the `start` button to begin streaming!**\n",
        "\n",
        "**Instructions**. Type in text prompts or upload an audio file to control the overall style of the generated music in real time. The sliders change the influence of each prompt on the overall output. The other controls change various aspects of the system behavior (expand below for additional information).\n",
        "\n",
        "⭐ **Finetuned model - extra features** ⭐ This demo includes additional prompt controls compared to the [main Magenta RT demo](https://github.com/magenta/magenta-realtime/blob/main/notebooks/Magenta_RT_Demo.ipynb) that allow you to steer the model towards styles that are more in-distribution with respect to the finetuning dataset. This can be beneficial when the finetuning data has a narrow coverage of sounds compared to pre-training (as is usually the case). Guiding the model towards in-distribution styles can turn the limited size and diversity of the finetuning data into a feature rather than a bug: out-of-domain prompts mixed with in-distribution embeddings often result in audio that retains some of the original prompt direction, while also being \"infused\" with sonic characteristics typical of the finetuning domain. Alongside the mean style embedding, we expose some cluster centroids as additional in-distribution prompts to achieve similar effects on more compact subspaces of the overall finetuning space.\n",
        "\n",
        "**Disclaimer**. Magenta RT's training data primarily consists of Western\n",
        "instrumental music. As a consequence, Magenta RT has incomplete coverage of both\n",
        "vocal performance and the broader landscape of rich musical traditions\n",
        "worldwide.\n",
        "\n",
        "<details>\n",
        "  <summary>Click to expand for additional information on the controls</summary>\n",
        "\n",
        "*   **extra_buffering_seconds**: Increase this value if you experience audio\n",
        "    drops during generation. This will come at the expense of a greater latency,\n",
        "    but might help with internet connection issues. *You need to relaunch the\n",
        "    cell if you choose to modify this value*.\n",
        "\n",
        "*   **sampling options**\n",
        "\n",
        "    *   **temperature**: This controls how *chaotic* the model behaves. Low\n",
        "        temperature values (e.g., 0.9) will make the model's choices more\n",
        "        predictable and stable. High values (e.g., 1.5) will encourage more\n",
        "        surprising and experimental musical ideas, but can also lead to\n",
        "        instability.\n",
        "\n",
        "    *   **topk**: This parameter filters the model's vocabulary at each step. It\n",
        "        forces the model to choose its next prediction only from the *k* most\n",
        "        likely options.\n",
        "\n",
        "        *   A **low `topk`** value (e.g., 40) restricts the model to a smaller,\n",
        "            safer palette of options. This leads to more coherent and\n",
        "            predictable music that is less likely to have dissonant errors, but\n",
        "            can sometimes feel repetitive.\n",
        "        *   A **high `topk`** value gives the model a much wider range of\n",
        "            choices, allowing for more variety and unexpected turns. This can\n",
        "            make the output more creative, but also noisier.\n",
        "\n",
        "    *   **guidance**: This controls how strictly the generated music should\n",
        "        adhere to the **text prompts**.\n",
        "\n",
        "        *   A **higher value** will push the model to produce a textbook example\n",
        "            of the chosen style, emphasizing its key characteristics.\n",
        "        *   A **lower value** will treat the text prompts more as a loose\n",
        "            inspiration, allowing the model more creative freedom to deviate and\n",
        "            blend other influences.\n",
        "\n",
        "*   **Reset**: stop audio, and resets the model.\n",
        "\n",
        "*   **In-distribution steering**: This allows you to steer the model towards the finetuning style distribution.\n",
        "    The first slider controls the weight of the embedding corresponding to the mean style prompt, the remaining five each correspond to cluster centroids.\n",
        "\n",
        "*   **Text prompts**: Next to each text prompt is a slider that controls how\n",
        "    much each prompt should be affecting the model. This allows the creation of\n",
        "    *mixed* embeddings (try mixing synthwave and flamenco guitar together !).\n",
        "    You can also type your own prompt and modify existing ones.\n",
        "\n",
        "*   **Audio prompts**: Instead of using text to define a musical style, you can\n",
        "    also upload audio references! Click on the `Upload audio file` button to\n",
        "    create a new audio-based prompt. Note that only **the first 10s** of audio\n",
        "    will be used. Supported formats include `.wav`, `.mp3` and `.ogg`.\n",
        "\n",
        "</details>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "bZ92jCA7qzRQ"
      },
      "outputs": [],
      "source": [
        "# @title **Run this cell** to load the selected checkpoint and start the demo\n",
        "\n",
        "import abc\n",
        "import concurrent.futures\n",
        "import functools\n",
        "from typing import Sequence\n",
        "\n",
        "import IPython.display as ipd\n",
        "import ipywidgets as ipw\n",
        "from magenta_rt import asset\n",
        "from magenta_rt import audio as audio_lib\n",
        "from magenta_rt import system\n",
        "from magenta_rt.colab import prompt_types\n",
        "from magenta_rt.colab import utils\n",
        "from magenta_rt.colab import widgets\n",
        "import numpy as np\n",
        "\n",
        "extra_buffering_seconds = 0  # @param {\"type\":\"slider\",\"min\":0,\"max\":4,\"step\":0.1}\n",
        "BUFFERING_AMOUNT_SAMPLES = int(np.ceil(extra_buffering_seconds * 48000))\n",
        "\n",
        "if checkpoint_to_load.value == \"holly_finetune\":\n",
        "  checkpoint_path = \"checkpoints/llm_large_holly_finetune.tar\"\n",
        "  checkpoint_dir = asset.fetch(checkpoint_path, is_dir=True, extract_archive=True, source=\"gcp\")\n",
        "  mean_style_embed_path = \"finetune_features/holly_mean_style_embedding.npy\"\n",
        "  cluster_centroids_path = \"finetune_features/holly_cluster_centroids.npy\"\n",
        "  mean_style_embed = np.load(asset.fetch(mean_style_embed_path, source=\"gcp\"))\n",
        "  cluster_centroids = np.load(asset.fetch(cluster_centroids_path, source=\"gcp\"))\n",
        "else:\n",
        "  checkpoint_dir = f\"{MODEL_OUTPUT_DIR}/checkpoint_{checkpoint_to_load.value}\"\n",
        "print(\"Loading checkpoint from \", checkpoint_dir)\n",
        "\n",
        "MRT = system.MagentaRT(\n",
        "    tag=\"large\" if checkpoint_to_load.value == \"holly_finetune\" else MODEL_SIZE,\n",
        "    lazy=False,\n",
        "    checkpoint_dir=checkpoint_dir,\n",
        ")\n",
        "\n",
        "\n",
        "class AudioStreamer(abc.ABC):\n",
        "  \"\"\"Audio streamer base class.\"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      sample_rate: int = 48000,\n",
        "      num_channels: int = 2,\n",
        "      buffer_size: int = 48000 * 2,\n",
        "      extra_buffering: int = BUFFERING_AMOUNT_SAMPLES,\n",
        "  ):\n",
        "    self.audio_streamer = None\n",
        "    self.sample_rate = sample_rate\n",
        "    self.num_channels = num_channels\n",
        "    self.buffer_size = buffer_size\n",
        "    self.extra_buffering = extra_buffering\n",
        "\n",
        "  def on_stream_start(self):\n",
        "    \"\"\"Called when the UI starts streaming.\"\"\"\n",
        "    if self.audio_streamer is not None:\n",
        "      self.audio_streamer.reset_ring_buffer()\n",
        "\n",
        "  def on_stream_stop(self):\n",
        "    \"\"\"Called when the UI stops streaming.\"\"\"\n",
        "    pass\n",
        "\n",
        "  @property\n",
        "  @abc.abstractmethod\n",
        "  def warmup(self) -> bool:\n",
        "    \"\"\"Returns whether to warm up the audio streamer.\"\"\"\n",
        "    pass\n",
        "\n",
        "  def reset(self):\n",
        "    if self.audio_streamer is not None:\n",
        "      self.audio_streamer.reset_ring_buffer()\n",
        "\n",
        "  def start(self):\n",
        "    self.audio_streamer = utils.AudioStreamer(\n",
        "        self,\n",
        "        rate=self.sample_rate,\n",
        "        buffer_size=self.buffer_size,\n",
        "        warmup=self.warmup,\n",
        "        num_output_channels=self.num_channels,\n",
        "        additional_buffered_samples=self.extra_buffering,\n",
        "        start_streaming_callback=self.on_stream_start,\n",
        "        stop_streaming_callback=self.on_stream_stop,\n",
        "    )\n",
        "    self.reset()\n",
        "\n",
        "  def stop(self):\n",
        "    if self.audio_streamer is not None:\n",
        "      del self.audio_streamer\n",
        "      self.audio_streamer = None\n",
        "\n",
        "  def global_ui_params(self):\n",
        "    return utils.Parameters.get_values()\n",
        "\n",
        "  def get_prompts(self):\n",
        "    params = self.global_ui_params()\n",
        "    num_prompts = sum(map(lambda s: \"prompt_value\" in s, params.keys()))\n",
        "    prompts = []\n",
        "    for i in range(num_prompts):\n",
        "      prompt_weight = params[f\"prompt_weight_{i}\"]\n",
        "      prompt_value = params[f\"prompt_value_{i}\"]\n",
        "\n",
        "      if prompt_value is None or not prompt_weight:\n",
        "        continue\n",
        "\n",
        "      match type(prompt_value):\n",
        "        case prompt_types.TextPrompt:\n",
        "          prompt_value = prompt_value.strip()\n",
        "        case prompt_types.AudioPrompt:\n",
        "          pass\n",
        "        case prompt_types.EmbeddingPrompt:\n",
        "          pass\n",
        "        case _:\n",
        "          raise ValueError(f\"Unsupported prompt type: {type(prompt_value)}\")\n",
        "\n",
        "      prompts.append((prompt_value, prompt_weight))\n",
        "    return prompts\n",
        "\n",
        "  @abc.abstractmethod\n",
        "  def generate(self, ui_params):\n",
        "    pass\n",
        "\n",
        "  def __call__(self, inputs):\n",
        "    del inputs\n",
        "    return self.generate(self.global_ui_params())\n",
        "\n",
        "\n",
        "class MagentaRTStreamer(AudioStreamer):\n",
        "  \"\"\"Audio streamer class for our open weights Magenta RT model.\n",
        "\n",
        "  This class holds a pretrained Magenta RT model, a generation state and an\n",
        "  asynchronous executor to handle the embedding of text prompt without\n",
        "  interrupting the audio thread.\n",
        "\n",
        "  Args:\n",
        "    system: A MagentaRTBase instance.\n",
        "    mean_style_embed: Mean embedding of the finetuning data for in-distribution\n",
        "      steering (optional).\n",
        "    cluster_centroids: Cluster centroids of the finetuning data embeddings for\n",
        "      in-distribution steering (optional).\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      system: system.MagentaRTBase,\n",
        "      mean_style_embed: np.ndarray | None = None,\n",
        "      cluster_centroids: Sequence[np.ndarray] | None = None,\n",
        "  ):\n",
        "    super().__init__()\n",
        "    self.system = system\n",
        "    self.state = None\n",
        "    self.executor = concurrent.futures.ThreadPoolExecutor()\n",
        "    self.mean_style_embed = mean_style_embed\n",
        "    self.cluster_centroids = cluster_centroids\n",
        "\n",
        "  @property\n",
        "  def warmup(self):\n",
        "    return True\n",
        "\n",
        "  @functools.cache\n",
        "  def embed_style(self, style: str):\n",
        "    return self.executor.submit(self.system.embed_style, style)\n",
        "\n",
        "  @functools.cache\n",
        "  def embed_audio(self, audio: tuple[float]):\n",
        "    audio = audio_lib.Waveform(np.asarray(audio), 16000)\n",
        "    return self.executor.submit(self.system.embed_style, audio)\n",
        "\n",
        "  def get_style_embedding(self, force_wait: bool = False):\n",
        "    prompts = self.get_prompts()\n",
        "    weighted_embedding = np.zeros((768,), dtype=np.float32)\n",
        "    total_weight = 0.0\n",
        "    for prompt_value, prompt_weight in prompts:\n",
        "      match type(prompt_value):\n",
        "        case prompt_types.TextPrompt:\n",
        "          if not prompt_value:\n",
        "            continue\n",
        "          embedding = self.embed_style(prompt_value)\n",
        "\n",
        "        case prompt_types.AudioPrompt:\n",
        "          embedding = self.embed_audio(tuple(prompt_value.value))\n",
        "\n",
        "        case prompt_types.EmbeddingPrompt:\n",
        "          embedding = prompt_value.value\n",
        "\n",
        "        case _:\n",
        "          raise ValueError(f\"Unsupported prompt type: {type(prompt_value)}\")\n",
        "\n",
        "      if isinstance(embedding, concurrent.futures.Future):\n",
        "        if force_wait:\n",
        "          embedding.result()\n",
        "\n",
        "        if not embedding.done():\n",
        "          continue\n",
        "\n",
        "        embedding = embedding.result()\n",
        "\n",
        "      weighted_embedding += embedding * prompt_weight\n",
        "      total_weight += prompt_weight\n",
        "\n",
        "    if self.mean_style_embed is not None:\n",
        "      params = self.global_ui_params()\n",
        "      training_weight = params[\"mean\"]\n",
        "      weighted_embedding += self.mean_style_embed * training_weight\n",
        "      total_weight += training_weight\n",
        "\n",
        "    if self.cluster_centroids is not None:\n",
        "      params = self.global_ui_params()\n",
        "      for i, centroid in enumerate(self.cluster_centroids):\n",
        "        centroid_weight = params[f\"centroid_{i}\"]\n",
        "        weighted_embedding += centroid * centroid_weight\n",
        "        total_weight += centroid_weight\n",
        "\n",
        "    if total_weight > 0:\n",
        "      weighted_embedding /= total_weight\n",
        "\n",
        "    return weighted_embedding\n",
        "\n",
        "  def on_stream_start(self):\n",
        "    self.get_style_embedding(force_wait=False)\n",
        "    self.get_style_embedding(force_wait=True)\n",
        "    super().on_stream_start()\n",
        "\n",
        "  def reset(self):\n",
        "    self.state = None\n",
        "    self.embed_style.cache_clear()\n",
        "    super().reset()\n",
        "\n",
        "  def generate(self, ui_params):\n",
        "    chunk, self.state = self.system.generate_chunk(\n",
        "        state=self.state,\n",
        "        style=self.get_style_embedding(),\n",
        "        seed=None,\n",
        "        **ui_params,\n",
        "    )\n",
        "    return chunk.samples\n",
        "\n",
        "  def stop(self):\n",
        "    self.executor.shutdown(wait=True)\n",
        "\n",
        "\n",
        "# BUILD UI\n",
        "\n",
        "\n",
        "def build_prompt_ui(default_prompts: Sequence[str], num_audio_prompt: int):\n",
        "  \"\"\"Add interactive prompt widgets and register them.\"\"\"\n",
        "  prompts = []\n",
        "\n",
        "  for p in default_prompts:\n",
        "    prompts.append(widgets.Prompt())\n",
        "    prompts[-1].text.value = p\n",
        "\n",
        "  prompts[0].slider.value = 1.0\n",
        "\n",
        "  # add audio prompt\n",
        "  for _ in range(num_audio_prompt):\n",
        "    prompts.append(widgets.AudioPrompt())\n",
        "    prompts[-1].slider.value = 0.0\n",
        "\n",
        "  utils.Parameters.register_ui_elements(\n",
        "      display=False,\n",
        "      **{f\"prompt_weight_{i}\": p.slider for i, p in enumerate(prompts)},\n",
        "      **{f\"prompt_value_{i}\": p.prompt_value for i, p in enumerate(prompts)},\n",
        "  )\n",
        "  return [p.get_widget() for p in prompts]\n",
        "\n",
        "\n",
        "def build_steering_option_ui(num_centroids):\n",
        "  \"\"\"Add interactive steering option widgets and register them.\"\"\"\n",
        "  options = {\n",
        "      \"mean\": ipw.FloatSlider(\n",
        "          min=0.0,\n",
        "          max=2.0,\n",
        "          step=0.01,\n",
        "          value=1.0,\n",
        "          description=\"mean\",\n",
        "          layout=ipw.Layout(width=\"500px\"),\n",
        "      ),\n",
        "  }\n",
        "\n",
        "  for i in range(num_centroids):\n",
        "    options[f\"centroid_{i}\"] = ipw.FloatSlider(\n",
        "        min=0.0,\n",
        "        max=2.0,\n",
        "        step=0.01,\n",
        "        value=0.0,\n",
        "        description=f\"centroid {i+1}\",\n",
        "        layout=ipw.Layout(width=\"500px\"),\n",
        "    )\n",
        "\n",
        "  utils.Parameters.register_ui_elements(display=False, **options)\n",
        "\n",
        "  return list(options.values())\n",
        "\n",
        "\n",
        "def build_sampling_option_ui():\n",
        "  \"\"\"Add interactive sampling option widgets and register them.\"\"\"\n",
        "  options = {\n",
        "      \"temperature\": ipw.FloatSlider(\n",
        "          min=0.0,\n",
        "          max=4.0,\n",
        "          step=0.01,\n",
        "          value=1.3,\n",
        "          description=\"temperature\",\n",
        "      ),\n",
        "      \"topk\": ipw.IntSlider(\n",
        "          min=0,\n",
        "          max=1024,\n",
        "          step=1,\n",
        "          value=40,\n",
        "          description=\"topk\",\n",
        "      ),\n",
        "      \"guidance_weight\": ipw.FloatSlider(\n",
        "          min=0.0,\n",
        "          max=10.0,\n",
        "          step=0.01,\n",
        "          value=5.0,\n",
        "          description=\"guidance\",\n",
        "      ),\n",
        "  }\n",
        "\n",
        "  utils.Parameters.register_ui_elements(display=False, **options)\n",
        "\n",
        "  return list(options.values())\n",
        "\n",
        "\n",
        "utils.Parameters.reset()\n",
        "\n",
        "# Initialize streamer\n",
        "streamer = MagentaRTStreamer(\n",
        "    MRT,\n",
        "    mean_style_embed=mean_style_embed,\n",
        "    cluster_centroids=cluster_centroids,\n",
        ")\n",
        "\n",
        "\n",
        "def _reset_state(*args, **kwargs):\n",
        "  del args, kwargs\n",
        "  streamer.reset()\n",
        "\n",
        "\n",
        "reset_button = ipw.Button(description=\"reset\")\n",
        "reset_button.on_click(_reset_state)\n",
        "\n",
        "\n",
        "# Building interactive UI\n",
        "ipd.display(\n",
        "    ipw.VBox([\n",
        "        widgets.area(\n",
        "            \"sampling options\",\n",
        "            *build_sampling_option_ui(),\n",
        "            reset_button,\n",
        "        ),\n",
        "        widgets.area(\n",
        "            \"in-distribution steering\",\n",
        "            *build_steering_option_ui(len(cluster_centroids)),\n",
        "        ),\n",
        "        widgets.area(\n",
        "            \"prompts\",\n",
        "            *build_prompt_ui(\n",
        "                [\n",
        "                    \"synthwave\",\n",
        "                    \"flamenco guitar\",\n",
        "                    \"\",\n",
        "                    \"\",\n",
        "                ],\n",
        "                num_audio_prompt=2,\n",
        "            ),\n",
        "        ),\n",
        "    ])\n",
        ")\n",
        "\n",
        "streamer.start()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xmvBk8UOphdL"
      },
      "source": [
        "# License and disclaimer\n",
        "\n",
        "Magenta RealTime is offered under a combination of licenses: the codebase is\n",
        "licensed under\n",
        "[Apache 2.0](https://github.com/magenta/magenta-realtime/blob/main/LICENSE),\n",
        "and the model weights under\n",
        "[Creative Commons Attribution 4.0 International](https://creativecommons.org/licenses/by/4.0/legalcode).\n",
        "\n",
        "In addition, we specify the following usage terms:\n",
        "\n",
        "Copyright 2025 Google LLC\n",
        "\n",
        "Use these materials responsibly and do not generate content, including outputs,\n",
        "that infringe or violate the rights of others, including rights in copyrighted\n",
        "content.\n",
        "\n",
        "Google claims no rights in outputs you generate using Magenta RealTime. You and\n",
        "your users are solely responsible for outputs and their subsequent uses.\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, all software and\n",
        "materials distributed here under the Apache 2.0 or CC-BY licenses are\n",
        "distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,\n",
        "either express or implied. See the licenses for the specific language governing\n",
        "permissions and limitations under those licenses. You are solely responsible for\n",
        "determining the appropriateness of using, reproducing, modifying, performing,\n",
        "displaying or distributing the software and materials, and any outputs, and\n",
        "assume any and all risks associated with your use or distribution of any of the\n",
        "software and materials, and any outputs, and your exercise of rights and\n",
        "permissions under the licenses."
      ]
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
